#!/usr/bin/env python
"""

This script automates the application/removal of iRules from VIPs

=====
Usage
=====

::

	irulemgr.py [options] <host> <user> <irule> <regex>

:param host: Bigip hostname
:param user: login username
:param irule: the iRule to add/remove
:param regex: regular expression used to match VIPs

-------
Options
-------

  --version     show program's version number and exit
  -h, --help    show this help message and exit
  -r, --remove  removes the iRule from the VIP(s) matched instead of adding
                them
  -l, --list    lists the iRules currently available on the Bigip, then exit
  -p PARTITION, --partition=PARTITION    set the name of the parition to work with (default=Common)

========
Examples
========

*Listing iRules on a Bigip*::

    $ irulemgr.py -l 172.16.248.10 admin
    Password: 

    172.16.248.10 has the following iRules:

    test1-irule
    test2-irule
    test3-irule

*Adding an iRule to all VIPs ending in '-80'*::

	$ irulemgr.py 172.16.248.10 admin test1-irule '\-80$'
	Password: 

	This will apply the "test1-irule" iRule on the following vips:

	vip2-80
	vip12-80
	vip3-80
	vip13-80

	Continue?(y/N): y
	iRule "test1-irule" was successfully applied to the "vip2-80" vip
	iRule "test1-irule" was successfully applied to the "vip12-80" vip
	iRule "test1-irule" was successfully applied to the "vip3-80" vip
	iRule "test1-irule" was successfully applied to the "vip13-80" vip
	There were no errors

*Removing an iRule from all VIPs ending in '-80'*::

	$ irulemgr.py -r 172.16.248.10 admin test1-irule '\-80$'
	Password: 

	This will remove the "test1-irule" iRule from the following vips:

	vip2-80
	vip12-80
	vip3-80
	vip13-80

	Continue?(y/N): y
	iRule "test1-irule" was successfully removed from the "vip2-80" vip
	iRule "test1-irule" was successfully removed from the "vip12-80" vip
	iRule "test1-irule" was successfully removed from the "vip3-80" vip
	iRule "test1-irule" was successfully removed from the "vip13-80" vip
	There were no errors

"""
import re,sys,getpass,os,logging,suds
sys.path.append("lib/python/")
import pycontrol.pycontrol as pc
from optparse import OptionParser
logging.basicConfig(level=logging.CRITICAL)

if pc.__version__ < '2.0':
    print "Requires pycontrol version 2.x!"
    sys.exit()

#### This section updated automatically by subversion, then used by optparse
svnrev = "$Rev$"
rev = re.search('^\$Rev: (\d+)', svnrev).group(1)
version = "1.0.%s" % rev
#### End section



def get_options_and_args():
	# Getting our processing of arguments and options out of the way
	usage = "usage: %prog [options] <host> <username> <iRule> <vip regex>"
	description = "This script automates the application/removal of iRules from VIPs"
	parser = OptionParser(usage=usage, description=description, version=version)
	parser.add_option("-r", "--remove", dest="remove", help="removes the iRule from the VIP(s) matched instead of adding them",
					action="store_true", default=False)
	parser.add_option("-l", "--list", dest="list", help="lists the iRules currently available on the Bigip, then exit",
					action="store_true", default=False)
	parser.add_option("-p", "--partition", dest="partition", help="set the name of the parition to work with (default=Common)",
					type="string", default="Common")
	# Get password from user
	(options, args) = parser.parse_args()
	upass = getpass.getpass()
	print ""
	return (options, args, upass)

def error_handling(f):
	if re.search('error_string         : (.*)', f.fault.faultstring):
		print "Error: " + re.search('error_string         : (.*)', f.fault.faultstring).group(1)
	else:
		print f.fault.faultstring
		
def cont():
	answer = raw_input("Continue?(y/N): ")
	if re.search('[yY][Ee]?[Ss]?', answer):
		pass
	elif re.search('[Nn][Oo]?', answer) or answer == "":
		print "\nOk, exiting..."
		sys.exit()
	else:
		print "\nI don't understand your answer: " + answer
		sys.exit()


def convert_rule(z, rule, priority):
	rule_seq = z.typefactory.create('LocalLB.VirtualServer.VirtualServerRuleSequence')
	rule_set = z.typefactory.create('LocalLB.VirtualServer.VirtualServerRule')
	rule_set.rule_name = rule
	rule_set.priority = priority
	rule_seq.item = [rule_set]
	return rule_seq

	
def remove_irule_from_vips(vip_names, irule_name, vs):
	problems = 0
	for vip_name in vip_names:
		try:
			vs.remove_rule(virtual_servers = [vip_name], rules = [convert_rule(vs,irule_name, 500)])
		except suds.WebFault, f:
			error_handling(f)
			problems += 1
		else:
			print "iRule \"%s\" was successfully removed from the \"%s\" vip" % (irule_name, vip_name)
	return problems
	
def apply_irule_to_vips(vip_names, irule_name, vs):
	problems = 0
	for vip_name in vip_names:
		priority = 500
		while True:
			try:
				vs.add_rule(virtual_servers = [vip_name], rules = [convert_rule(vs,irule_name, priority)])
			except suds.WebFault, f:
				if re.search('priorities.*unique', f.fault.faultstring):
					priority += 1
					continue
				else:
					error_handling(f)
					problems += 1
					break
			else:
				print "iRule \"%s\" was successfully applied to the \"%s\" vip" % (irule_name, vip_name)
				break
	return problems
	
def list_irules(ir):
	all_irules = ir.get_list()
	# Alphabetically sort the irule list
	all_irules.sort()
	irules = []
	for irule in all_irules:
		# Hide the built-in sys iRules
		if re.search('^_', irule):
			pass
		else:
			irules.append(irule)
	return irules
	
def main(options, args, upass):
	
	if len(args) == 2 and options.list:
		(hostname, username) = args
	elif len(args) != 4:
		parser.error("incorrect number of arguments")
	else:
		(hostname, username, irule_name, pattern) = args
		
	bigip = pc.BIGIP(
	        hostname = hostname,
	        username = username,
	        password = upass,
	        fromurl = True,
	        wsdls = ['LocalLB.VirtualServer', 'LocalLB.Rule', 'Management.Partition'])

	vs = bigip.LocalLB.VirtualServer
	ir = bigip.LocalLB.Rule
	
	try:
	    bigip.Management.Partition.set_active_partition(options.partition)
	except suds.WebFault, f:
		error_handling(f)
		sys.exit(1) 
	except:
		print "Error: Unable to connect to the F5.  Is the bigip reachable and username/password correct?"
		sys.exit(1)
	else:
		if options.list:
			print hostname + " has the following iRules:\n"
			irules = list_irules(ir)
			for irule in irules:
				print irule
			print ""
			sys.exit()
		else:	
			vip_names = []
			all_vips = vs.get_list()
			for vip in all_vips:
				if re.search(pattern, vip):
					vip_names.append(vip)
			if options.remove:
				print "This will remove the \"%s\" iRule from the following vips:\n" % irule_name
				for vip in vip_names:
					print vip
				print ""
				cont()
				problems = remove_irule_from_vips(vip_names, irule_name, vs)
			else:
				print "This will apply the \"%s\" iRule on the following vips:\n" % irule_name
				for vip in vip_names:
					print vip
				print ""
				cont()
				problems = apply_irule_to_vips(vip_names, irule_name, vs)
			
			if problems == 0:
				print "There were no errors"
			elif problems == 1:
				print "There was 1 error"
			else:
				print "There were %s errors" % problems

if __name__ == "__main__":
	(options, args, upass) = get_options_and_args()
	main(options, args, upass)
